# encoding: utf-8
import torch

from constant import *
from models import *
from utils.images import batch_tensor_to_img

netG = UnetGenerator(input_nc, output_nc, 8, ngf).to(device)

weight_path = r"E:\PythonMyCode\cv_about\watermark-cleaner-demo\cGAN_pix2pix\weights\netG_epoch_39_err_0.53367_l1_6.19942.pth"

netG.load_state_dict(
    torch.load(weight_path, map_location=device)
)
netG.to(device)

img_path = r"E:\PythonMyCode\cv_about\watermark-cleaner-demo\utils\imgs\origin\YK2383_6.jpg"

inference_dataset = DataGen4(file_path=img_path)
inference_dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=workers)

with torch.no_grad():
    for img, label in inference_dataloader:
        image = img.to(device)
        predict = netG(image)
        epoch = 0
        batch_tensor_to_img(epoch, predict)



